#include "OpenCL.h"

#include <GL/glx.h>
#include <iostream>
#include <fstream>
#include <cmath>

#include "Particle.h"
#include "Body.h"

extern Body** bodies;
extern Particle** particles;
extern int totalNumberOfBodies;
extern int totalNumberOfParticles;
extern GLuint bodyVBO;
extern GLuint particleVBO;
extern float timeDelta;
extern float particleRadius;
extern int vertexCount;
extern float terminalVelocity;
extern float worldSize;
extern float springCoefficient;
extern float dampingCoefficient;

cl::BufferGL bufferBodyVBO;
cl::BufferGL bufferParticleVBO;
cl::Context context;
cl::CommandQueue queue;

cl::Kernel updateGridKernel;
cl::Kernel resetGridKernel;
cl::Kernel collisionDetectionKernel;
cl::Kernel updateMomentaKernel;
cl::Kernel performStepKernel;
cl::Kernel updateParticlesKernel;

cl::Buffer bufferBodyMass;
cl::Buffer bufferBodyForce;
cl::Buffer bufferBodyPosition;
cl::Buffer bufferBodyVelocity;
cl::Buffer bufferBodyLinearMomentum;
cl::Buffer bufferBodyQuaternion;
cl::Buffer bufferBodyRotationMatrix;
cl::Buffer bufferAngularVelocity;
cl::Buffer bufferAngularMomentum;
cl::Buffer bufferInitialIITDiagonal;
cl::Buffer bufferInverseInertiaTensor;

cl::Buffer bufferParticleMass;
cl::Buffer bufferParticlePosition;
cl::Buffer bufferParticleVelocity;
cl::Buffer bufferParticleForce;

cl_int* clCountGrid;
cl_int4* clIndexGrid;
cl_int3* clParticleGridIndex;

cl::Buffer bufferCountGrid;
cl::Buffer bufferIndexGrid;
cl::Buffer bufferParticleGridIndex;

cl::NDRange particlesGlobal;
cl::NDRange particlesLocal;
cl::NDRange bodiesGlobal;
cl::NDRange bodiesLocal;
cl::NDRange gridGlobal;
cl::NDRange gridLocal;


extern float voxelSize;
extern int gridLength;
extern int gridSize;

cl_float3 gridMinimumPositionVector;

void initOpenCLGrid() {

	voxelSize = 2.0f*particleRadius;
	gridLength = (int)ceil(2.0f * worldSize / voxelSize);
	/* 2 extra grid voxels in each dimension in case particles go slightly outside the world borders */
	gridLength += 2; 

	gridSize = gridLength * gridLength * gridLength;

	std::cout << "Grid dimensions: " << gridLength << "*" << gridLength << "*" << gridLength << std::endl;
	std::cout << "Grid voxel size: " << voxelSize << std::endl;

	gridMinimumPositionVector.s[0] = -worldSize - voxelSize;
	gridMinimumPositionVector.s[1] = -voxelSize;
	gridMinimumPositionVector.s[2] = -worldSize - voxelSize;

	std::cout << "Grid minimum position: (" << gridMinimumPositionVector.s[0] << ", " <<
	       gridMinimumPositionVector.s[1] << ", " << gridMinimumPositionVector.s[2] << ")" << std::endl;
	std::cout << "Grid maximum position: (" << gridMinimumPositionVector.s[0] + (gridLength*voxelSize) << ", " <<
		gridMinimumPositionVector.s[1] + (gridLength*voxelSize) << ", " <<
		gridMinimumPositionVector.s[2] + (gridLength*voxelSize) << ")" << std::endl;

	clCountGrid = new cl_int[gridSize];
	clIndexGrid = new cl_int4[gridSize];
	clParticleGridIndex = new cl_int3[totalNumberOfParticles];
	
	for (int i=0; i<gridSize; i++) {
		clCountGrid[i] = 0;
		clIndexGrid[i].s[0] = -1;
		clIndexGrid[i].s[1] = -1;
		clIndexGrid[i].s[2] = -1;
		clIndexGrid[i].s[3] = -1;
	}

	bufferCountGrid = cl::Buffer(context, CL_MEM_READ_WRITE, gridSize * sizeof(cl_int));
	bufferIndexGrid = cl::Buffer(context, CL_MEM_READ_WRITE, gridSize * sizeof(cl_int4));
	bufferParticleGridIndex = cl::Buffer(context, CL_MEM_READ_WRITE, totalNumberOfParticles * sizeof(cl_int3));

	queue.enqueueWriteBuffer(bufferCountGrid, CL_TRUE, 0,
			gridSize * sizeof(cl_int), clCountGrid);

	queue.enqueueWriteBuffer(bufferIndexGrid, CL_TRUE, 0,
			gridSize * sizeof(cl_int4), clIndexGrid);
}

void initCL() {

	cl::vector<cl::Platform> platforms;
	cl::Platform::get(&platforms);

	GLXContext glCtx = glXGetCurrentContext();
	cl_context_properties cpsGL[] = { CL_CONTEXT_PLATFORM, (cl_context_properties)platforms[0](),
		CL_GLX_DISPLAY_KHR, (intptr_t)glXGetCurrentDisplay(),
		CL_GL_CONTEXT_KHR, (intptr_t)glCtx, 0};

	context = cl::Context(CL_DEVICE_TYPE_GPU, cpsGL);

	cl::vector<cl::Device> devices = context.getInfo<CL_CONTEXT_DEVICES>();
	
	std::cout << "Device name: " << devices[0].getInfo<CL_DEVICE_NAME>().c_str() << std::endl;
	std::cout << "Device vendor: " << devices[0].getInfo<CL_DEVICE_VENDOR>().c_str() << std::endl;

	queue = cl::CommandQueue(context, devices[0]);

	std::ifstream sourceFile("rigidbody_kernels.cl");
        std::string sourceCode(
            std::istreambuf_iterator<char>(sourceFile),
            (std::istreambuf_iterator<char>()));
	cl::Program::Sources source(1, std::make_pair(sourceCode.c_str(), sourceCode.length()+1));
	cl::Program program = cl::Program(context, source);
	program.build(devices);

	updateGridKernel = cl::Kernel(program, "updateGrid");
	resetGridKernel = cl::Kernel(program, "resetGrid");
	collisionDetectionKernel = cl::Kernel(program, "collisionDetection");
	updateMomentaKernel = cl::Kernel(program, "updateMomenta");
	performStepKernel = cl::Kernel(program, "performStep");
	updateParticlesKernel = cl::Kernel(program, "updateParticles");

	//Create OpenCL buffers from VBOs
	bufferBodyVBO = cl::BufferGL(context, CL_MEM_WRITE_ONLY, bodyVBO);
	bufferParticleVBO = cl::BufferGL(context, CL_MEM_WRITE_ONLY, particleVBO);

	initializeHostDataArrays();	
	initOpenCLGrid();
	initializeDeviceBuffers();

	updateHostDataArrays();
	writeHostDataToDevice();

	updateGridKernel.setArg(0, bufferCountGrid);
	updateGridKernel.setArg(1, bufferIndexGrid);
	updateGridKernel.setArg(2, bufferParticlePosition);
	updateGridKernel.setArg(3, gridMinimumPositionVector);
	updateGridKernel.setArg(4, voxelSize);
	updateGridKernel.setArg(5, gridLength);
	updateGridKernel.setArg(6, bufferParticleGridIndex);

	resetGridKernel.setArg(0, bufferCountGrid);
	resetGridKernel.setArg(1, bufferIndexGrid);

	collisionDetectionKernel.setArg(0, bufferParticleMass);
	collisionDetectionKernel.setArg(1, bufferParticlePosition);
	collisionDetectionKernel.setArg(2, bufferParticleVelocity);
	collisionDetectionKernel.setArg(3, bufferParticleForce);
	collisionDetectionKernel.setArg(4, particleRadius);
	collisionDetectionKernel.setArg(5, worldSize);
	collisionDetectionKernel.setArg(6, springCoefficient);
	collisionDetectionKernel.setArg(7, dampingCoefficient);
	collisionDetectionKernel.setArg(8, bufferParticleGridIndex);
	collisionDetectionKernel.setArg(9, bufferCountGrid);
	collisionDetectionKernel.setArg(10, bufferIndexGrid);
	collisionDetectionKernel.setArg(11, gridLength);

	updateMomentaKernel.setArg(0, bufferBodyMass);
	updateMomentaKernel.setArg(1, bufferBodyForce);
	updateMomentaKernel.setArg(2, bufferBodyPosition);
	updateMomentaKernel.setArg(3, bufferBodyLinearMomentum);
	updateMomentaKernel.setArg(4, bufferAngularMomentum);
	updateMomentaKernel.setArg(5, bufferParticlePosition);
	updateMomentaKernel.setArg(6, bufferParticleForce);
	updateMomentaKernel.setArg(7, timeDelta);
	updateMomentaKernel.setArg(8, terminalVelocity);

	performStepKernel.setArg(0, bufferBodyMass);
	performStepKernel.setArg(1, bufferBodyPosition);
	performStepKernel.setArg(2, bufferBodyVelocity);
	performStepKernel.setArg(3, bufferBodyLinearMomentum);
	performStepKernel.setArg(4, bufferBodyQuaternion);
	performStepKernel.setArg(5, bufferBodyRotationMatrix);
	performStepKernel.setArg(6, bufferAngularVelocity);
	performStepKernel.setArg(7, bufferAngularMomentum);
	performStepKernel.setArg(8, bufferInitialIITDiagonal);
	performStepKernel.setArg(9, bufferInverseInertiaTensor);
	performStepKernel.setArg(10, bufferBodyVBO);
	performStepKernel.setArg(11, timeDelta);
	performStepKernel.setArg(12, particleRadius);
	int bodyVBOStride = vertexCount * 3;
	performStepKernel.setArg(13, bodyVBOStride);

	updateParticlesKernel.setArg(0, bufferBodyPosition);
	updateParticlesKernel.setArg(1, bufferBodyVelocity);
	updateParticlesKernel.setArg(2, bufferBodyRotationMatrix);
	updateParticlesKernel.setArg(3, bufferAngularVelocity);
	updateParticlesKernel.setArg(4, bufferParticleVBO);
	updateParticlesKernel.setArg(5, bufferParticlePosition);
	updateParticlesKernel.setArg(6, bufferParticleVelocity);
	updateParticlesKernel.setArg(7, particleRadius);

	particlesGlobal = cl::NDRange(totalNumberOfParticles);
	particlesLocal = cl::NDRange(150);
	bodiesGlobal = cl::NDRange(totalNumberOfBodies);
	bodiesLocal = cl::NDRange(100);
	gridGlobal = cl::NDRange(gridSize);
	gridLocal = cl::NDRange(gridLength);

}

void initializeHostDataArrays() {
	clBodyMass = new cl_float[totalNumberOfBodies];
	clBodyForce = new cl_float3[totalNumberOfBodies];
	clBodyPosition = new cl_float3[totalNumberOfBodies];
	clBodyVelocity = new cl_float3[totalNumberOfBodies];
	clBodyLinearMomentum = new cl_float3[totalNumberOfBodies];
	clBodyQuaternion = new cl_float4[totalNumberOfBodies];
	clBodyRotationMatrix = new cl_float[9*totalNumberOfBodies];
	clBodyAngularVelocity = new cl_float3[totalNumberOfBodies];
	clBodyAngularMomentum = new cl_float3[totalNumberOfBodies];
	clInitialIITDiagonal = new cl_float3[totalNumberOfBodies];
	clBodyInverseInertiaTensor = new cl_float[9*totalNumberOfBodies];

	clParticleMass = new cl_float[totalNumberOfParticles];
	clParticlePosition = new cl_float3[totalNumberOfParticles];
	clParticleVelocity = new cl_float3[totalNumberOfParticles];
	clParticleForce = new cl_float3[totalNumberOfParticles];
}

void initializeDeviceBuffers() {

	bufferBodyMass = cl::Buffer(context, CL_MEM_READ_ONLY, 
			totalNumberOfBodies * sizeof(cl_float));

	bufferBodyForce = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float3));
	
	bufferBodyPosition = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float3));

	bufferBodyVelocity = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float3));

	bufferBodyLinearMomentum = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float3));

	bufferBodyQuaternion = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float4));

	bufferBodyRotationMatrix = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float) * 9);

	bufferAngularVelocity = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float3));

	bufferAngularMomentum = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float3));

	bufferInitialIITDiagonal = cl::Buffer(context, CL_MEM_READ_ONLY, 
			totalNumberOfBodies * sizeof(cl_float3));

	bufferInverseInertiaTensor = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfBodies * sizeof(cl_float) * 9);


	bufferParticleMass = cl::Buffer(context, CL_MEM_READ_ONLY, 
			totalNumberOfParticles * sizeof(cl_float));

	bufferParticlePosition = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfParticles * sizeof(cl_float3));

	bufferParticleVelocity = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfParticles * sizeof(cl_float3));

	bufferParticleForce = cl::Buffer(context, CL_MEM_READ_WRITE, 
			totalNumberOfParticles * sizeof(cl_float3));

}

void updateHostDataArrays() {

	for (int i=0; i<totalNumberOfBodies; i++) {
		bodies[i]->updateOpenCLArray(i);
	}
	for (int i=0; i<totalNumberOfParticles; i++) {
		particles[i]->updateOpenCLArray(i);
	}
}

void writeHostDataToDevice() {

	queue.enqueueWriteBuffer(bufferBodyMass, CL_TRUE, 0, 
			totalNumberOfBodies * sizeof(cl_float), clBodyMass);
	queue.enqueueWriteBuffer(bufferBodyForce, CL_TRUE, 0, 
			totalNumberOfBodies * sizeof(cl_float3), clBodyForce);
	queue.enqueueWriteBuffer(bufferBodyPosition, CL_TRUE, 0, 
			totalNumberOfBodies * sizeof(cl_float3), clBodyPosition);
	queue.enqueueWriteBuffer(bufferBodyVelocity, CL_TRUE, 0, 
			totalNumberOfBodies * sizeof(cl_float3), clBodyVelocity);
	queue.enqueueWriteBuffer(bufferBodyLinearMomentum, CL_TRUE, 0, 
			totalNumberOfBodies * sizeof(cl_float3), clBodyLinearMomentum);
	queue.enqueueWriteBuffer(bufferBodyQuaternion, CL_TRUE, 0,
			totalNumberOfBodies * sizeof(cl_float4), clBodyQuaternion);
	queue.enqueueWriteBuffer(bufferBodyRotationMatrix, CL_TRUE, 0,
			totalNumberOfBodies * sizeof(cl_float) * 9, clBodyRotationMatrix);
	queue.enqueueWriteBuffer(bufferAngularVelocity, CL_TRUE, 0,
			totalNumberOfBodies * sizeof(cl_float3), clBodyAngularVelocity);
	queue.enqueueWriteBuffer(bufferAngularMomentum, CL_TRUE, 0,
			totalNumberOfBodies * sizeof(cl_float3), clBodyAngularMomentum);
	queue.enqueueWriteBuffer(bufferInitialIITDiagonal, CL_TRUE, 0,
			totalNumberOfBodies * sizeof(cl_float3), clInitialIITDiagonal);
	queue.enqueueWriteBuffer(bufferInverseInertiaTensor, CL_TRUE, 0,
			totalNumberOfBodies * sizeof(cl_float) * 9, clBodyInverseInertiaTensor);

	queue.enqueueWriteBuffer(bufferParticleMass, CL_TRUE, 0,
			totalNumberOfParticles * sizeof(cl_float), clParticleMass);
	queue.enqueueWriteBuffer(bufferParticlePosition, CL_TRUE, 0,
			totalNumberOfParticles * sizeof(cl_float3), clParticlePosition);
	queue.enqueueWriteBuffer(bufferParticleVelocity, CL_TRUE, 0,
			totalNumberOfParticles * sizeof(cl_float3), clParticleVelocity);
	queue.enqueueWriteBuffer(bufferParticleForce, CL_TRUE, 0,
			totalNumberOfParticles * sizeof(cl_float3), clParticleForce);

}


void stepSimulationWithOpenCL() {
	
	cl::vector<cl::Memory> vboBuffers;
	vboBuffers.push_back(bufferParticleVBO);
	vboBuffers.push_back(bufferBodyVBO);

	queue.enqueueAcquireGLObjects(&vboBuffers);

	queue.enqueueNDRangeKernel(resetGridKernel, cl::NullRange, gridGlobal, gridLocal);
	queue.enqueueNDRangeKernel(updateGridKernel, cl::NullRange, particlesGlobal, particlesLocal);
	queue.enqueueNDRangeKernel(collisionDetectionKernel, cl::NullRange, particlesGlobal, particlesLocal);
	queue.enqueueNDRangeKernel(updateMomentaKernel, cl::NullRange, bodiesGlobal, bodiesLocal);
	queue.enqueueNDRangeKernel(performStepKernel, cl::NullRange, bodiesGlobal, bodiesLocal);
	queue.enqueueNDRangeKernel(updateParticlesKernel, cl::NullRange, particlesGlobal, particlesLocal);

	queue.enqueueReleaseGLObjects(&vboBuffers);
	queue.finish();
	
}

